import bitarray
import collections
import heapq

class huffman_tree():

    def __init__(self, symbol, weight, left=None, right=None):
        self.symbol = symbol
        self.weight = weight
        self.left = left
        self.right = right

    def __lt__(self, right):
        return self.weight < right.weight

    def symbol_to_bitarray(self, symbol):
        if symbol == self.symbol:
            return bitarray.bitarray()
        elif self.left is not None and symbol in self.left.symbol:
            return bitarray.bitarray([0]) + \
                self.left.symbol_to_bitarray(symbol)
        elif self.right is not None and symbol in self.right.symbol:
            return bitarray.bitarray([1]) + \
                self.right.symbol_to_bitarray(symbol)
        else:
            raise ValueError("symbol %s not in tree" % symbol)


def huffman(text):
    if len(text) == 0:
        return {}
    if len(text) == 1:
        return {text: bitarray.bitarray([0])}

    W = collections.Counter(text)

    trees = []
    for symbol, weight in W.items():
        trees.append(huffman_tree(symbol, weight))

    heapq.heapify(trees)

    while len(trees) > 1:
        t1 = heapq.heappop(trees)
        t2 = heapq.heappop(trees)
        symbol = t1.symbol + t2.symbol
        weight = t1.weight + t2.weight
        tm = huffman_tree(symbol, weight, t1, t2)
        heapq.heappush(trees, tm)


    return {c: trees[0].symbol_to_bitarray(c) for c in W.keys()}
